{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"OMP_PROC_BIND\"] = os.environ.get(\"OMP_PROC_BIND\", \"true\")\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from pix_transform.pix_transform import PixTransform\n",
    "from baselines.baselines import bicubic\n",
    "from utils.utils import downsample,align_images\n",
    "from prox_tv import tvgen\n",
    "from utils.plots import plot_result\n",
    "\n",
    "####  load dataset  #############################################################\n",
    "data_path = \"./data/depth_sample_images.npz\"\n",
    "\n",
    "dataset = np.load(data_path)\n",
    "target_imgs = dataset[\"target_imgs\"].squeeze()\n",
    "guide_imgs =  dataset[\"guide_imgs\"].squeeze()\n",
    "dataset.close()\n",
    "\n",
    "####  define parameters  ########################################################\n",
    "params = {'img_idxs' : [], # idx images to process, if empty then all of them\n",
    "            \n",
    "          'scaling': 8,\n",
    "          'greyscale': False, # Turn image into grey-scale\n",
    "          'channels': -1,\n",
    "          \n",
    "          'spatial_features_input': True,\n",
    "          'weights_regularizer': [0.0001, 0.001, 0.0001], # spatial color head\n",
    "          'loss': 'l1',\n",
    " \n",
    "          'optim': 'adam',\n",
    "          'lr': 0.001,\n",
    "                  \n",
    "          'batch_size': 32,\n",
    "          'iteration': 1024*32*32//32,\n",
    "                  \n",
    "          'logstep': 64,\n",
    "          \n",
    "          'final_TGV' : False, # Total Generalized Variation in post-processing\n",
    "          'align': False, # Move image around for evaluation in case guide image and target image are not perfectly aligned\n",
    "          'delta_PBP': 1, # Delta for percentage of bad pixels \n",
    "         }\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if len(params['img_idxs'])==0:\n",
    "    idxs = np.array(range(0,target_imgs.shape[0]))\n",
    "else:\n",
    "    idxs = params['img_idxs']\n",
    "\n",
    "for n_image,idx in enumerate(idxs):\n",
    "    \n",
    "    print(\"####### image {}/{} - image idx {} ########\".format(n_image+1,len(idxs),idx))\n",
    "    \n",
    "    guide_img = guide_imgs[idx]\n",
    "    target_img = target_imgs[idx]\n",
    "    source_img = downsample(target_img,params['scaling'])\n",
    "\n",
    "    bicubic_target_img = bicubic(source_img=source_img, scaling_factor=params['scaling'])\n",
    "    \n",
    "    predicted_target_img = PixTransform(guide_img=guide_img,source_img=source_img,params=params,target_img=target_img)\n",
    "    \n",
    "\n",
    "\n",
    "    if params['final_TGV'] :\n",
    "        print(\"applying TGV...\")\n",
    "        predicted_target_img = tvgen(predicted_target_img,[0.1, 0.1],[1, 2],[1, 1])\n",
    "        \n",
    "    if params['align'] :\n",
    "        print(\"aligning...\")\n",
    "        target_img,predicted_target_img = align_images(target_img,predicted_target_img)\n",
    "\n",
    "    \n",
    "    f, ax = plot_result(guide_img,source_img,predicted_target_img,bicubic_target_img,target_img)\n",
    "    plt.show()\n",
    "    \n",
    "    if target_img is not None:\n",
    "        # compute metrics and plot results\n",
    "        MSE = np.mean((predicted_target_img - target_img) ** 2)\n",
    "        MAE = np.mean(np.abs(predicted_target_img - target_img))\n",
    "        PBP = np.mean(np.abs(predicted_target_img - target_img) > params[\"delta_PBP\"])\n",
    "\n",
    "        print(\"MSE: {:.3f}  ---  MAE: {:.3f}  ---  PBP: {:.3f}\".format(MSE,MAE,PBP))\n",
    "        print(\"\\n\\n\")\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:myenv]",
   "language": "python",
   "name": "conda-env-myenv-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
